package com.gearbox.platform.slurm.service;

import com.gearbox.core.model.shell.ShellResult;
import com.gearbox.core.util.CmdExecuteResultAssertion;
import com.gearbox.core.util.ShellExecutor;

import java.util.List;
import java.util.Locale;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

public class SlurmCommandExecutor {

    private static final String GET_NODE_INFO_COMMAND = "sinfo -N -p %s";

    private static final String DELETE_NODE_COMMAND = "scontrol delete NodeName=%s";

    private static final String DELETE_HOSTS_RECORD = "sed -i '/%s/d' /etc/hosts";

    private static final String UPDATE_NODE_STATUS = "scontrol update NodeName=%s State=%s Reason=%s";

    private static final String LIST_JOB_BY_STATE = "squeue -o \"%%i %%j %%u %%C %%m %%T\" --state=%s";

    private static final String LIST_JOB_BY_NODE = "squeue -w %s";

    private static final String WAITING_JOB_STATE = "PD";

    private static final String RUNNING_JOB_STATE = "R";

    private static final String LIST_ACC_COMMAND = "sacctmgr list ass";

    private static final String UPDATE_REASON = "update by Gearbox";

    private static final String REASON_STRING = "\"%s %s\"";

    private static final Pattern NODE_QUOTA_PATTERN = Pattern.compile("node=[0-9]+");

    private static final Pattern NODE_QUOTA_NUMBER_PATTERN = Pattern.compile("[0-9]+");

    public static List<String> getNodesInfoByPartition(String partition) {
        String cmd = String.format(Locale.ROOT, GET_NODE_INFO_COMMAND, partition);
        ShellResult result = ShellExecutor.runCommand(cmd);
        CmdExecuteResultAssertion.assertError(result);
        return result.getReturnValues()
            .stream()
            .filter(Pattern.compile(partition).asPredicate())
            .collect(Collectors.toList());
    }

    public static void deleteNodeByName(String name) {
        String cmd = String.format(Locale.ROOT, DELETE_NODE_COMMAND, name);
        ShellResult result = ShellExecutor.runCommand(cmd);
        CmdExecuteResultAssertion.assertError(result);
    }

    public static void deleteHostsRecordByKeyword(String keyword) {
        String cmd = String.format(Locale.ROOT, DELETE_HOSTS_RECORD, keyword);
        ShellResult result = ShellExecutor.runCommand(cmd);
        CmdExecuteResultAssertion.assertError(result);
    }

    public static void updateNodeStatus(String nodeName, String status, String reason) {
        String reasonString = String.format(Locale.ROOT, REASON_STRING, reason, UPDATE_REASON);
        String cmd = String.format(Locale.ROOT, UPDATE_NODE_STATUS, nodeName, status, reasonString);
        ShellResult result = ShellExecutor.runCommand(cmd);
        CmdExecuteResultAssertion.assertError(result);
    }

    public static List<String> listWaitingJobInfos() {
        String cmd = String.format(Locale.ROOT, LIST_JOB_BY_STATE, WAITING_JOB_STATE);
        ShellResult result = ShellExecutor.runCommand(cmd);
        CmdExecuteResultAssertion.assertError(result);
        return result.getReturnValues();
    }

    public static List<String> listRunningJobInfos() {
        String cmd = String.format(Locale.ROOT, LIST_JOB_BY_STATE, RUNNING_JOB_STATE);
        ShellResult result = ShellExecutor.runCommand(cmd);
        CmdExecuteResultAssertion.assertError(result);
        return result.getReturnValues();
    }

    public static List<String> getQuotasByUser(String userName) {
        ShellResult result = ShellExecutor.runCommand(LIST_ACC_COMMAND);
        CmdExecuteResultAssertion.assertError(result);
        return result.getReturnValues()
            .stream()
            .filter(Pattern.compile(userName).asPredicate())
            .map(NODE_QUOTA_PATTERN::matcher)
            .filter(Matcher::find)
            .map(Matcher::group)
            .map(NODE_QUOTA_NUMBER_PATTERN::matcher)
            .filter(Matcher::find)
            .map(Matcher::group)
            .collect(Collectors.toList());
    }

    public static List<String> getJobInfosByNodeName(String nodeName) {
        String cmd = String.format(Locale.ROOT, LIST_JOB_BY_NODE, nodeName);
        ShellResult result =  ShellExecutor.runCommand(cmd).withCommandLine(cmd);
        CmdExecuteResultAssertion.assertError(result);
        return result.getReturnValues();
    }
}
